{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from typing import Union\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.append('..')\n",
    "\n",
    "from settings import data_home_folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 只包含 “氢”， (87744, 2)\n",
    "> 计划增加关键词"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd  \n",
    "from typing import Union, List  \n",
    "\n",
    "tqdm.pandas()\n",
    "def filter_by_keyword_optimized(  \n",
    "    df_: pd.DataFrame,\n",
    "    cols: Union[List[str], str], \n",
    "    keyword: Union[List[str], str]\n",
    ") -> pd.DataFrame:\n",
    "    if isinstance(cols, str):\n",
    "        cols = [cols]\n",
    "    if isinstance(keyword, str):  \n",
    "        keyword = [keyword]\n",
    "\n",
    "    # 使用正则表达式模式匹配所有关键字（可选，但仅当关键字数量很大时可能更有效）  \n",
    "    # pattern = '|'.join(r'\\b{}\\b'.format(re.escape(kw)) for kw in keyword)  \n",
    "\n",
    "    # 直接在选定的列上应用str.contains，并使用any沿行方向聚合结果  \n",
    "    mask = df_[cols].progress_apply(\n",
    "            lambda row: row.astype(str).str.contains('|'.join(keyword), na=False).any(), \n",
    "            axis=1\n",
    "        )  \n",
    "    return df_[mask]\n",
    "\n",
    "# 示例用法  \n",
    "# df = pd.DataFrame({...})  # 您的DataFrame  \n",
    "# result = filter_by_keyword_optimized(df, 'column_name', 'keyword')  \n",
    "# 或者使用多个列和多个关键字  \n",
    "# result = filter_by_keyword_optimized(df, ['col1', 'col2'], ['kw1', 'kw2'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 679/679 [4:52:45<00:00, 25.87s/it]   \n"
     ]
    }
   ],
   "source": [
    "# 遍历 data_home_folder 下的所有文件，找出csv文件\n",
    "series = []\n",
    "for file_name in tqdm(os.listdir(data_home_folder)):\n",
    "    if file_name.endswith(\".csv\"):\n",
    "        tmp_df = pd.read_csv(\n",
    "            os.path.join(data_home_folder, file_name),\n",
    "            # usecols=[\"企业名称\", \"经营范围\"],\n",
    "            low_memory=False,\n",
    "        )\n",
    "        series.append(filter_by_keyword_optimized(\n",
    "                tmp_df, \n",
    "                [\"企业名称\", \"经营范围\"], \n",
    "                [\"氢\", \"膜\", \"气\", \"燃料电池\", \"LNG\", \"CNG\", \"电解\", \"电堆\"]\n",
    "            ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "339"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(series)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_df = pd.concat(series)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(87744, 2)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8675611, 32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_df.to_csv(\"氢能_名称_经营范围_867w.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_df.head(1).to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 二次筛选"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 初次筛选出的数据量太大了，想进行二次筛选"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_169583/775312148.py:1: DtypeWarning: Columns (10,16,18,22,31) have mixed types. Specify dtype option on import or set low_memory=False.\n",
      "  df = pd.read_csv(\"氢能_名称_经营范围_867w.csv\")\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(\"氢能_名称_经营范围_867w.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 8675611/8675611 [09:30<00:00, 15204.79it/s]\n"
     ]
    }
   ],
   "source": [
    "# 遍历 data_home_folder 下的所有文件，找出csv文件\n",
    "res2 = filter_by_keyword_optimized(\n",
    "        df,\n",
    "        [\"企业名称\", \"经营范围\"], \n",
    "        [\"氢\", \"膜\", \"燃料电池\", \"LNG\", \"CNG\", \"电解\", \"电堆\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1616747, 32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res2.head(1).to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1616747/1616747 [01:48<00:00, 14867.73it/s]\n"
     ]
    }
   ],
   "source": [
    "res3 = filter_by_keyword_optimized(\n",
    "        res2,\n",
    "        [\"企业名称\", \"经营范围\"], \n",
    "        [\"氢\", \"燃料电池\", \"LNG\", \"CNG\", \"电解\", \"电堆\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(112328, 32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "res3.to_csv(\"氢能_名称_经营范围_11w.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
